from .priority_queue_base import PriorityQueueBase
import math


class HeapPriorityQueue(PriorityQueueBase):
    """
    A min-oriented priority queue implemented with a binary heap
    使用数组二叉树实现堆（小顶堆）
    """

    def __init__(self, contents=()):
        """Create empty priority queue"""
        self._data = [self._Item(k, v) for k, v in contents]
        if len(self._data) > 1:
            self._heapify()

    def _heapify(self):
        start = self._parent(len(self) - 1) # 从最后叶子的parent开始
        for j in range(start, -1, -1):
            self._downheap(j)

    def __len__(self):
        return len(self._data)

    def __str__(self):
        ret = ""
        for i, item in enumerate(self._data):
            ret += str(item._key)
            if i == 2**int(math.log(i+1, 2) + 1) -2 or i == len(self._data) - 1:
                ret += '\n'
            else:
                ret += ', '
        return ret

    def _parent(self, j):
        return (j - 1) // 2

    def _left(self, j):
        return 2 * j + 1

    def _right(self, j):
        return 2 * j + 2

    def _has_right(self, j):
        return self._right(j) < len(self._data)

    def _has_left(self, j):
        return self._left(j) < len(self._data)

    def swap(self, i, j):
        """swap the element i and element j"""
        self._data[i], self._data[j] = self._data[j], self._data[i]

    def _upheap(self, j):
        """插入后向上冒泡"""
        parent = self._parent(j)
        if j > 0 and self._data[j] < self._data[parent]:
            self.swap(j, parent)
            return self._upheap(parent)

    def _downheap(self, j):
        """
        删除后向下冒泡
        1.如果j没有孩子，即算法终止
        2.如果j没有右孩子，则令c表示j的左孩子
        3.否则，比较j的两个孩子，选择较小权值的孩子
        """
        if self._has_left(j):
            left = self._left(j)
            small_child = left
            if self._has_right(j):
                right = self._right(j)
                if self._data[right] < self._data[left]:
                    small_child = right
            if self._data[small_child] < self._data[j]:
                self.swap(j, small_child)
                self._downheap(small_child)

    def add(self, key, value):
        """Add a key-value pair to the priority queue"""
        self._data.append(self._Item(key, value))
        self._upheap(len(self._data) - 1)

    def min(self):
        """Return but not remove minimum key-value"""
        if self.is_empty():
            raise Exception("heap priority queue is empty")
        min_item = self._data[0]
        return (min_item._key, min_item._value)

    def remove_min(self):
        """
        Remove and return the minimum key-value
        O(logn)
        """
        if self.is_empty():
            raise Exception("heap priority queue is empty")
        self.swap(0, len(self._data)-1)  # put minimum item at the end
        min_item = self._data.pop()  # delete minimun item
        self._downheap(0)  # fix new root
        return (min_item._key, min_item._value)
